#!/usr/bin/env python

#############################################################################
# network.py: Module for parsing and sending network commands.
#############################################################################
# Copyright (C) 2012
#
# Authors: Dan Larkin <danielhlarkin@gmail.com>
#          Yonatan Naamad <cached@gmail.com>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#############################################################################

import struct
import logging

#central server communication
JOIN_op, JOIN_schema = 0, ("!16s I I H H", lambda dct: ' 32s'*dct['n_current'], 
                             'k','b_total','d','n_max','n_current','L')
ALLOCATE_op, ALLOCATE_schema = 1, ("!I I 20s", None, 'b_min', 'b_max', 'T')
REQUEST_op, REQUEST_schema = 2, ("!I 20s", None, 'b', 'T')
FREE_op, FREE_schema = 3, ("!I 20s", None, 'b_free', 'T')

#torrent server communication
GET_op, GET_schema = 6, ("!20s", None, 'T')
GIVE_op, GIVE_schema = 7, ("!I", lambda dct: str(dct['l']) + 's', 'l', 'F')

def send_packet(socket, op, schema, d):
    '''Sends a packet given its schema and corresponding value dictionary 'd'.'''
    data = write_schema(d, schema)
    socket.send(chr(op) + data)
    logging.info("Sent packet with opcode %d." % (op,))

def read_packet(socket, op, schema):
    '''Read packet from the given schema into a dictionary.'''
    #figure out size of constant part of packet
    osize = struct.calcsize(schema[0])
    
    #get that from the socket
    try:
        p1 = socket.recv(osize)
    except Exception as e:
        logging.error("Failed to read '%d' bytes from socket." % osize)
        raise e
        return None
        
    if schema[1] == None: 
        #parse the constant portion
        data = read_schema(p1,schema)
        return data

    #parse the constant portion
    data1 = read_schema(p1,(schema[0], None) + schema[2:-1])

    #figure out size of the rest of the packet
    tsize = struct.calcsize(schema[1](data1))
    
    #get rest of packet from socket
    try:
        p2 = ""
        j = 1
        while j != 0:
            q = socket.recv(tsize)
            p2 += q
            j = len(q)
    except Exception as e:
        logging.error("Failed to read '%d' bytes from socket for rest of packet." % p2)
        raise e
        return None
        
    #reparse the entire schema
    data2 = read_schema(p1 + p2, schema)
    
    return data2
    
def read_schema(packet, schema):
    '''Takes a packet and extracts info from it according to the given schema.
       Schema entries are given in some tuple, in the form (format, f_dep?, *p)
       where format is the format of the constant-sized portion of the input
       as in the module struct, f_dep? is either None or a function generating
       the rest (nonconstant portion) of the schema based on the constant part,
       and *p is the list of names of variables we are assigning. If there is a
       nonconstant part, it is returned as a list of objects.'''
    d = {}
    
    unp1 = struct.unpack(schema[0],packet[:struct.calcsize(schema[0])])
    
    for k,v in zip(schema[2:],unp1):
        d[k] = v
        
    
    if schema[1] != None:
        new_schema = schema[0] + " " + schema[1](d)
        unp2 = struct.unpack(new_schema, packet)
        
        d = {}
        d[schema[-1]] = []
        for k,v in zip(list(schema[2:]) + ([schema[-1]]*len(unp2)),unp2):
            if k == schema[-1]:
                d[k].append(v)
            else:
                d[k] = v
    
    return d
    
def write_schema(d, schema):
    '''Inverts read_schema'''
    if schema[1] == None:
        new_schema = schema[0]
        params = tuple(map(lambda x: d[x], schema[2:]))
    else:
        new_schema = schema[0] + " " + schema[1](d)
        params = tuple(map(lambda x: d[x], schema[2:-1]) + d[schema[-1]])

    return struct.pack(new_schema, *params)
